# coding: utf-8
# definitions of TSP
import os

import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman')

line_colors = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9']
city_colors = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9']


class TSP:
    def __init__(self, path):
        self.path = path
        self.file_dir, self.name = os.path.split(path)  # instance name
        self.name = self.name.split('.')[0]
        self.data = np.loadtxt(path, skiprows=6, encoding='utf-8')
        self.data = np.insert(self.data, obj=0, values=0, axis=0)  # raw tsp data
        self.cities = self.data[1:, 0].astype(int).tolist()  # city indexes
        self.depot = self.cities[0]  # depot
        self.n = len(self.cities)  # number of cities
        self.x, self.y = self.data[:, 1].astype(float), self.data[:, 2].astype(float)  # coordinates
        self.distance_matrix = np.sqrt(np.square(self.x - self.x[np.arange(self.n + 1).reshape(self.n + 1, 1)])
                                       + np.square(self.y - self.y[np.arange(self.n + 1).reshape(self.n + 1, 1)]))
        self.distance_matrix[0, :], self.distance_matrix[:, 0] = 0, 0  # distance matrix
    
    def __repr__(self):
        return f'{self.__class__.__name__}:\n path={self.path}\n name={self.name}\n cities={self.cities}'
    
    def cost(self, route: list) -> float:
        return sum([self.distance_matrix[route[i - 1]][route[i]] for i in range(1, len(route))])
    
    def draw(self, plot, route=None, dot_size=None, rotate=False, show_text=False, label=False):
        dot_size = round(np.log2(self.n)) * 1 if dot_size is None else dot_size
        # dot_size = 0
        x, y = (self.y, self.x) if rotate else (self.x, self.y)
        plot.xticks([])
        plot.yticks([])
        plot.scatter(x[1:], y[1:], s=dot_size, c='grey')
        if route is not None:
            plot.plot(x[route], y[route])
        if show_text:
            for data in self.data[1:]:
                plot.text(data[1], data[2], s=data[0].astype(int), size=12)
        if label:
            plot.legend(loc=9, bbox_to_anchor=(0.5, 1.1), borderaxespad=0, ncol=4)
        return plot


# data = TSP(path=f'../../data/TSPLIB/att48.tsp.txt')
# print(data)
# print(data.distance_matrix)
# data.draw(plt, routes=[[1, 2, 3, 4, 5, 1], [1, 11, 12, 13, 1]])
# plt.show()
